#!/usr/bin/env python3
"""
数据集迁移工具
将现有的安全帽检测数据集扩展为综合安全检测数据集
"""

import os
import shutil
from pathlib import Path
import yaml

class DatasetMigrator:
    def __init__(self):
        """初始化数据集迁移器"""
        
        # 原始类别映射（安全帽检测）
        self.original_classes = {
            0: 'Helmet',
            1: 'No-Helmet', 
            2: 'Person'
        }
        
        # 新的类别映射（综合安全检测）
        self.new_classes = {
            0: 'Person',           # 人员
            1: 'Helmet',           # 佩戴安全帽
            2: 'No-Helmet',        # 未佩戴安全帽
            3: 'Safety-Vest',      # 佩戴安全背心/工服
            4: 'No-Safety-Vest',   # 未佩戴安全背心/工服
            5: 'Smoking',          # 抽烟行为
            6: 'Phone-Call',       # 接打电话
            7: 'Fire',             # 火焰
            8: 'Smoke',            # 烟雾
            9: 'Danger-Zone',      # 危险区域
            10: 'Safety-Equipment', # 安全设备
            11: 'Vehicle'          # 工程车辆
        }
        
        # 类别ID映射关系
        self.class_mapping = {
            0: 0,  # Person -> Person
            1: 1,  # Helmet -> Helmet  
            2: 2   # No-Helmet -> No-Helmet
        }
    
    def migrate_labels(self, source_dir, target_dir):
        """
        迁移标签文件，更新类别ID映射
        
        Args:
            source_dir: 原始标签目录
            target_dir: 目标标签目录
        """
        os.makedirs(target_dir, exist_ok=True)
        
        label_files = list(Path(source_dir).glob('*.txt'))
        print(f"发现 {len(label_files)} 个标签文件")
        
        for label_file in label_files:
            source_path = label_file
            target_path = Path(target_dir) / label_file.name
            
            # 读取原始标签
            with open(source_path, 'r') as f:
                lines = f.readlines()
            
            # 转换标签
            new_lines = []
            for line in lines:
                line = line.strip()
                if line:
                    parts = line.split()
                    if len(parts) >= 5:
                        old_class_id = int(parts[0])
                        
                        # 映射到新的类别ID
                        if old_class_id in self.class_mapping:
                            new_class_id = self.class_mapping[old_class_id]
                            new_line = f"{new_class_id} {' '.join(parts[1:])}\n"
                            new_lines.append(new_line)
                        else:
                            print(f"警告: 未知类别ID {old_class_id} 在文件 {label_file.name}")
            
            # 保存新标签
            with open(target_path, 'w') as f:
                f.writelines(new_lines)
        
        print(f"标签迁移完成: {source_dir} -> {target_dir}")
    
    def copy_images(self, source_dir, target_dir):
        """
        复制图片文件
        
        Args:
            source_dir: 原始图片目录
            target_dir: 目标图片目录
        """
        os.makedirs(target_dir, exist_ok=True)
        
        # 支持的图片格式
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']
        
        image_files = []
        for ext in image_extensions:
            image_files.extend(list(Path(source_dir).glob(f'*{ext}')))
            image_files.extend(list(Path(source_dir).glob(f'*{ext.upper()}')))
        
        print(f"发现 {len(image_files)} 个图片文件")
        
        for image_file in image_files:
            source_path = image_file
            target_path = Path(target_dir) / image_file.name
            
            # 复制文件
            shutil.copy2(source_path, target_path)
        
        print(f"图片复制完成: {source_dir} -> {target_dir}")
    
    def migrate_dataset(self, source_base_dir="./", target_base_dir="comprehensive_dataset"):
        """
        迁移整个数据集
        
        Args:
            source_base_dir: 原始数据集根目录
            target_base_dir: 目标数据集根目录
        """
        print("="*50)
        print("开始数据集迁移")
        print("="*50)
        
        # 定义数据集分割
        splits = ['train', 'valid', 'test']
        
        for split in splits:
            print(f"\n处理 {split} 数据集...")
            
            # 源路径
            source_images = os.path.join(source_base_dir, split, 'images')
            source_labels = os.path.join(source_base_dir, split, 'labels')
            
            # 目标路径
            target_images = os.path.join(target_base_dir, split, 'images')
            target_labels = os.path.join(target_base_dir, split, 'labels')
            
            # 检查源目录是否存在
            if os.path.exists(source_images):
                self.copy_images(source_images, target_images)
            else:
                print(f"警告: 源图片目录不存在 {source_images}")
            
            if os.path.exists(source_labels):
                self.migrate_labels(source_labels, target_labels)
            else:
                print(f"警告: 源标签目录不存在 {source_labels}")
        
        # 创建新的数据配置文件
        self.create_comprehensive_config(target_base_dir)
        
        print("\n" + "="*50)
        print("数据集迁移完成!")
        print(f"迁移后的数据集位于: {target_base_dir}")
        print("="*50)
    
    def create_comprehensive_config(self, base_dir):
        """创建综合数据集配置文件"""
        config = {
            'train': f'{base_dir}/train/images',
            'val': f'{base_dir}/valid/images',
            'test': f'{base_dir}/test/images',
            'nc': 12,
            'names': [
                'Person',           # 0
                'Helmet',           # 1
                'No-Helmet',        # 2
                'Safety-Vest',      # 3
                'No-Safety-Vest',   # 4
                'Smoking',          # 5
                'Phone-Call',       # 6
                'Fire',             # 7
                'Smoke',            # 8
                'Danger-Zone',      # 9
                'Safety-Equipment', # 10
                'Vehicle'           # 11
            ]
        }
        
        config_path = 'comprehensive_data.yaml'
        with open(config_path, 'w', encoding='utf-8') as f:
            yaml.dump(config, f, default_flow_style=False, allow_unicode=True)
        
        print(f"已创建配置文件: {config_path}")
    
    def analyze_dataset(self, dataset_dir):
        """分析数据集统计信息"""
        print("\n" + "="*50)
        print("数据集分析")
        print("="*50)
        
        splits = ['train', 'valid', 'test']
        total_stats = {class_id: 0 for class_id in self.new_classes.keys()}
        
        for split in splits:
            print(f"\n{split.upper()} 数据集:")
            
            labels_dir = os.path.join(dataset_dir, split, 'labels')
            images_dir = os.path.join(dataset_dir, split, 'images')
            
            if not os.path.exists(labels_dir):
                print(f"  标签目录不存在: {labels_dir}")
                continue
            
            if not os.path.exists(images_dir):
                print(f"  图片目录不存在: {images_dir}")
                continue
            
            # 统计图片数量
            image_files = []
            for ext in ['.jpg', '.jpeg', '.png', '.bmp']:
                image_files.extend(list(Path(images_dir).glob(f'*{ext}')))
                image_files.extend(list(Path(images_dir).glob(f'*{ext.upper()}')))
            
            print(f"  图片数量: {len(image_files)}")
            
            # 统计标签
            label_files = list(Path(labels_dir).glob('*.txt'))
            print(f"  标签文件数量: {len(label_files)}")
            
            # 统计各类别数量
            class_counts = {class_id: 0 for class_id in self.new_classes.keys()}
            
            for label_file in label_files:
                with open(label_file, 'r') as f:
                    for line in f:
                        line = line.strip()
                        if line:
                            parts = line.split()
                            if len(parts) >= 5:
                                class_id = int(parts[0])
                                if class_id in class_counts:
                                    class_counts[class_id] += 1
                                    total_stats[class_id] += 1
            
            # 显示类别统计
            print(f"  类别统计:")
            for class_id, count in class_counts.items():
                if count > 0:
                    class_name = self.new_classes[class_id]
                    print(f"    {class_name}: {count}")
        
        # 显示总体统计
        print(f"\n总体统计:")
        for class_id, count in total_stats.items():
            if count > 0:
                class_name = self.new_classes[class_id]
                print(f"  {class_name}: {count}")
    
    def create_training_plan(self):
        """创建训练计划建议"""
        plan = """
        
        ==========================================
        综合安全检测训练计划建议
        ==========================================
        
        1. 当前数据状态:
           ✓ 已有安全帽检测数据（Person, Helmet, No-Helmet）
           ⚠ 缺少其他安全类别数据
        
        2. 数据收集优先级:
           高优先级：
           - Safety-Vest / No-Safety-Vest (工服检测)
           - Smoking (抽烟检测)
           - Phone-Call (接打电话检测)
           
           中优先级：
           - Fire / Smoke (火焰烟雾检测)
           - Danger-Zone (危险区域)
           
           低优先级：
           - Safety-Equipment (安全设备)
           - Vehicle (车辆)
        
        3. 训练策略:
           阶段1: 使用现有数据训练基础模型
           阶段2: 逐步添加新类别数据，进行增量训练
           阶段3: 全类别联合训练优化
        
        4. 数据增强建议:
           - 针对建筑场景的特定增强
           - 光照变化模拟
           - 角度和距离变化
           - 遮挡情况模拟
        
        5. 评估指标:
           - 各类别单独mAP
           - 安全关键类别的召回率
           - 误检率控制
        
        ==========================================
        """
        
        print(plan)
        
        with open('训练计划建议.txt', 'w', encoding='utf-8') as f:
            f.write(plan)
        
        print("训练计划已保存到: 训练计划建议.txt")

def main():
    """主函数"""
    print("="*60)
    print("建筑安全检测数据集迁移工具")
    print("="*60)
    
    migrator = DatasetMigrator()
    
    # 检查原始数据集
    print("检查原始数据集...")
    original_data_exists = False
    
    for split in ['train', 'valid', 'test']:
        images_path = os.path.join(split, 'images')
        labels_path = os.path.join(split, 'labels')
        
        if os.path.exists(images_path) and os.path.exists(labels_path):
            original_data_exists = True
            print(f"✓ 找到 {split} 数据集")
        else:
            print(f"✗ 未找到 {split} 数据集")
    
    if not original_data_exists:
        print("\n未找到原始安全帽检测数据集")
        print("请确保数据集结构如下:")
        print("├── train/")
        print("│   ├── images/")
        print("│   └── labels/")
        print("├── valid/")
        print("│   ├── images/")
        print("│   └── labels/")
        print("└── test/")
        print("    ├── images/")
        print("    └── labels/")
        return
    
    # 询问是否进行迁移
    print("\n" + "="*50)
    user_input = input("是否开始数据集迁移? (y/n): ").lower().strip()
    
    if user_input in ['y', 'yes']:
        # 执行迁移
        migrator.migrate_dataset()
        
        # 分析迁移后的数据集
        migrator.analyze_dataset('comprehensive_dataset')
        
        # 生成训练计划
        migrator.create_training_plan()
        
        print("\n" + "="*50)
        print("下一步操作:")
        print("1. 运行 train_comprehensive_safety.py 开始训练")
        print("2. 或者继续添加其他类别的数据")
        print("3. 参考 训练计划建议.txt 制定数据收集计划")
        print("="*50)
    
    else:
        print("取消迁移操作")

if __name__ == "__main__":
    main() 